package com.xiaoqinag.ds.tree;

import com.alibaba.fastjson.JSONObject;

import java.util.TreeMap;

/**
 * @author 小强
 * @createTime 2018/10/29 19:37
 */
public class RBT {

    private Node root;

    public RBT() {
    }

    public RBT(Node root) {
        this.root = root;
    }

    public Node getTree() {
        return root;
    }

    public void insert(Node node) {
        insertIntoTree(root, node);
        fixAfterInsertion(node);
    }

    public void deleteNode(Node node) {
        if (node == null) {
            return;
        }
        if (node.getLeft() == null && node.getRight() == null) { //子节点均为空
            if (node == root) {
                root = null;
            } else {
                Node p = node.getParent();
                if (p == p.getLeft())
                    p.setLeft(null);
                if (p == node.getRight())
                    p.setRight(null);
            }
        } else if (node.getLeft() == null || node.getRight() == null) {    //一支为空
            Node replacement = node.getRight() == null ? node.getLeft() : node.getRight();
            if (node == node.getParent().getRight())
                node.getParent().setRight(replacement);
            if (node == node.getParent().getLeft())
                node.getParent().setLeft(replacement);
            node = null;
        } else {    //左右子节点均不为空
            Node s = successor(node);
            node.setData(s.getData());
            
        }
    }

    public void insertIntoTree(Node root, Node node) {
        if (root == null) {
            this.root = node;
            return;
        }
        if (root.compareTo(node) < 0) {
            if (root.getRight() == null)
                root.setRight(node);
            else
                insertIntoTree(root.getRight(), node);
        } else if (root.compareTo(node) > 0) {
            if (root.getLeft() == null)
                root.setLeft(node);
            else
                insertIntoTree(root.getLeft(), node);
        }
    }

    public void fixAfterInsertion(Node x) {
        while (x != null && x != root && x.getParent().isRed()) {
            if (parentOf(x) == parentOf(parentOf(x)).getLeft()) {
                Node y = rightOf(parentOf(parentOf(x)));
                if (y == null || !y.isRed()) {
                    if (x == rightOf(parentOf(x))) {
                        Node p = parentOf(x);
                        rotateLeft(x);
                        x = p;
                    }
                    setColor(parentOf(x), Node.NodeColor.Black);
                    setColor(parentOf(parentOf(x)), Node.NodeColor.Red);
                    rotateRight(parentOf(x));
                } else {
                    setColor(parentOf(x), Node.NodeColor.Black);
                    setColor(y, Node.NodeColor.Black);
                    setColor(parentOf(parentOf(x)), Node.NodeColor.Red);
                    x = parentOf(parentOf(x));
                }
            } else {
                Node y = leftOf(parentOf(parentOf(x)));
                if (y == null || !y.isRed()) {
                    if (x == leftOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateRight(x);
                    }
                    setColor(parentOf(x), Node.NodeColor.Black);
                    setColor(parentOf(parentOf(x)), Node.NodeColor.Red);
                    rotateLeft(parentOf(x));
                } else {
                    setColor(parentOf(x), Node.NodeColor.Black);
                    setColor(y, Node.NodeColor.Black);
                    setColor(parentOf(parentOf(x)), Node.NodeColor.Red);
                    x = parentOf(parentOf(x));
                }
            }
        }
        this.root.setColor(Node.NodeColor.Black);
    }


    public void fixAfterDeletion(Node p) {

    }

    public void rotateLeft(Node x) {
        if (x == null || x == root) {
            return;
        }
        Node y = parentOf(x);
        Node gp = parentOf(y);
        if (gp != null && gp.getLeft() == y) {
            gp.setLeft(x);
        } else if (gp != null && gp.getRight() == y) {
            gp.setRight(x);
        } else {
            root = x;
        }
        y.setRight(leftOf(x));
        x.setLeft(y);
    }

    public void rotateRight(Node x) {
        if (x == null) {
            return;
        }
        Node y = parentOf(x);
        Node gp = parentOf(y);
        if (gp != null && leftOf(parentOf(y)) == y) {
            gp.setLeft(x);
        } else if (gp != null && rightOf(parentOf(y)) == y) {
            gp.setRight(x);
        } else {
            root = x;
        }
        y.setLeft(rightOf(x));
        x.setRight(y);
    }

    public Node successor(Node p) {
        if (p == null) {
            return null;
        }
        Node r = p.getRight();
        while (r != null && r.getLeft() != null) {
            r = r.getLeft();
        }
        return r;
    }

    public Node predecessor(Node p) {
        if (p == null) {
            return null;
        }
        Node l = p.getLeft();
        while (l != null && l.getRight() != null) {
            l = l.getRight();
        }
        return l;
    }

    public static Node parentOf(Node x) {
        return x == null ? null : x.getParent();
    }

    public static Node leftOf(Node x) {
        return x == null ? null : x.getLeft();
    }

    public static Node rightOf(Node x) {
        return x == null ? null : x.getRight();
    }

    public static void setColor(Node x, Node.NodeColor color) {
        if (x != null) {
            x.setColor(color);
        }
    }

}
